#!/usr/bin/env python
#
# testing for keep-alive connection
#   test tablet-server performance
#

import time,socket,base64,fcntl,errno,sys
import threading
import Queue as queue

import celtuce.net.ioloop as ioloop
import celtuce.net.iostream as iostream
import celtuce.help.codec as codec

from celtuce.help import logging
from celtuce.helper import getTimestamp
from celtuce.event import Event,OP_INSERT,OP_ERASE,OP_FIND, OP_TEST_ECHO
from celtuce.response import Response

#######################################################################

class   Client(object) :
    def __init__(self, server_host, server_port) :
        self._server_host = server_host
        self._server_port = server_port

        self._END_TAG = "\r\n\r\n"

        self._socket = None
        self._stream = None

        self._keep_alive = True

        self._id_base = id(self)
        self._id_serial = 1
        self._id_maximum = 1024 * 1024 * 1024

        self._logger = logging.getLogger('client')
        ## __init__()

    def __del__(self) :
        self.__disconnect(force=True)
        ## __del__()

    ####-----------------------------------------------------------####

    def __connect(self) :
        if self._socket : 
            #self._stream = iostream.IOStream( self._socket, self._io_loop )
            return

        try :
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            s.connect( (self._server_host, self._server_port) )
        except :
            self._logger.excpt()
            raise
        self._socket = s

        self._io_loop = ioloop.IOLoop()
        self._stream = iostream.IOStream( self._socket, self._io_loop )
        self._stream.set_close_callback( self.onDisconnect )

        return
        ## __connect()

    def __disconnect(self, force=False) :
        if self._keep_alive and not force :
            return

        if self._stream :
            self._stream.close()
            self._stream = None

        if self._socket :
            try :
                self._socket.close()
            except :
                self._logger.excpt()

        self._socket = None
        self._io_loop = None
        ## __disconnect()

    def onDisconnect(self) :
        self._logger.info('socket closed.')
        self._io_loop.stop()
        self.__disconnect(force=True)
        ## onDisconnect()

    ####-----------------------------------------------------------####

    def rpc(self, request) :
        self.__connect()

        #self._logger.debug('write request : %s', request)
        self._stream.write( ''.join((request, self._END_TAG)), self.__onWriteComplete )
        self._stream.read_until(self._END_TAG, self.__onResponse)
        self._response = None

        self._io_loop.start()

        self.__disconnect()
        return self._response
        ## rpc()

    def __onWriteComplete(self) :
        #self._logger.debug('write done.')
        pass
        ## __onWriteComplete()

    def __onResponse(self, response_str) :
        #self._logger.debug('read response : %s', repr(response_str))
        response = response_str.strip()
        if response :
            self._response = response_str.strip()
            self._io_loop.stop()
        else :
            self._stream.read_until(self._END_TAG, self.__onResponse)
        ## __onResponse()

    ####-----------------------------------------------------------####

    def genRequestID(self) :
        # TODO: lock
        serial = self._id_serial

        self._id_serial += 1
        if self._id_serial >= self._id_maximum :
            self._id_serial = 1 

        id = '.'.join((str(self._id_base), str(serial)))
        return id
        ## genRequestID()

    ## class Client


#######################################################################


def echo(any_data) :
    return any_data.strip()
    ## echo()


TCP_PORT = 14310

def runTest() :
    logger = logging.getLogger('test')

    client = Client('', TCP_PORT)

    table_name = 'test.gate4'

    try :
        # ping
        key = 'hello'
        data = 'world'
        timestamp_ms = getTimestamp()

        request = Event()
        request.op = OP_TEST_ECHO
        request.body = (table_name, key, data, timestamp_ms)
        request.request_id = client.genRequestID()
        request_str = codec.encode(request)

        response_str = client.rpc( request_str )
        response = codec.decode( response_str )

        logger.info('REQUEST  : %s', request)
        logger.info('RESPONSE : %s %s', response, response.body)

        # test codec performance
        time_start = time.time()
        times = 10000
        for i in xrange(times) :
            encoded  = codec.encode( request )
            response = codec.decode( encoded )
        time_last = time.time() - time_start    
        tps = times / time_last
        logger.info('test codec %d times last %.2f seconds, tps=%.2f', times, time_last, tps)

        # test echo
        key_prefix = 'hello'
        data_prefix = 'world'
        time_start = time.time()
        times = 1000
        sum_bytes = 0
        count_success = 0
        count_failure = 0
        for i in range(times) :
            key = '-'.join((key_prefix, str(i)))
            data = '-'.join((data_prefix, str(i)))
            timestamp_ms = getTimestamp()

            request = Event()
            request.op = OP_TEST_ECHO
            request.body = (table_name, key, str(timestamp_ms))
            request.request_id = client.genRequestID()
            request_str = codec.encode(request)

            response_str = client.rpc( request_str )
            response = codec.decode( response_str )
            if isinstance(response.body, dict) :
                response.body = tuple([ response.body[i] for i in range(len(response.body)) ])

            if response.body and request.body == response.body :
                count_success += 1
            else :
                print 'GOT %r, but WANT %r' % (response.body, request.body)
                count_failure += 1

            sum_bytes += len(response_str)
        time_last = time.time() - time_start

        tps = times / time_last
        bw = sum_bytes / time_last / 1024.0
        logger.info('rpc-FIND %d times last %.2f seconds, tps:%.2f, bw:%.2f KB/s count(success/failure)=%d/%d',times,time_last,tps,bw,count_success,count_failure)

    finally :
        logger.info('shutdown ...')

    ## runTest()


########################################################################

if  __name__ == '__main__' :

    runTest()


